import math
import torch
from torch.nn.modules.module import Module
from torch.nn.parameter import Parameter

'''
   图卷积神经网络
   输出向量x1（n，in_features）->输出向量x2（n，out_features）
   计算公式
   x2 = D-1/2  @  A  @  D-1/2   @   x1    @    W + b
    D-1/2：（n，n）
    A：邻接矩阵（n,n）
    x1 ：（n，in_features）
    W : 参数（in_features，out_features）
    b ：偏置（n，out_features）
    使用：
    x = torch.rand([100,4,5])
    adj = torch.rand([4,4])
    model1 = GCN(in_features = 2,out_features = 1)
    y = model1(x,adj)

'''


class GCN(Module):  #输入为（节点数，特征数（in_features））  输出为（节点数，特征数（out_features））  dim = 0为batch_size
    def __init__(self,in_features,out_features,bias = True):
        super(GCN,self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features,out_features)).type(torch.float32)  #初始化参数 维度:(in_features, out_features)
        if bias:
            self.bias =  Parameter(torch.FloatTensor(out_features)).type(torch.float32)  #初始化参数 维度:(out_features)
        else:
            self.register_parameter('bias',None)
        self.reset_parameters()
    def reset_parameters(self):  #初始化权重
        stdv = 1./math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv,stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv,stdv)
    def cal_D_(self,adj):  #计算度矩阵的-1/2次方
        # 计算度矩阵 D
        degree_matrix = torch.diag(adj.sum(dim=1))
        # 计算度矩阵的 -1/2 次方
        degree_matrix_sqrt = torch.sqrt(torch.linalg.pinv(degree_matrix))
        return degree_matrix_sqrt
        #输入数据  x(节点数，特征数（in_features）)    临界矩阵adj(节点数,节点数)
    def forward(self,x,adj):
        support = torch.matmul(x,self.weight.type(torch.float32))
        output = torch.bmm(adj.unsqueeze(0).expand(support.size(0),*adj.size()),support)
        if self.bias is not None:
            return output + self.bias.type(torch.float32)
        else:
            return output   
    def __repr__(self):
        return self.__class__.__name__ + ' (' \
            + str(self.in_features) + ' -> ' \
            + str(self.out_features) + ')'